#!/bin/env python
# -*- coding: utf-8 -*-
import taichi as ti
import numpy as np
from tqdm import tqdm
from scipy.spatial import cKDTree
from typing import List
from ..dither_util import DitherUtil
import cv2

@ti.data_oriented
class AtkinsonDitherer(DitherUtil):
    """
    Atkinson抖动算法实现类。
    支持灰度和彩色图像。
    """
    
    def __init__(self, arch=ti.cpu, debug=False, max_value=255):
        """
        初始化Atkinson抖动处理器。
        
        :param arch: Taichi架构，默认为CPU，可以设置为ti.gpu启用GPU加速
        :param debug: 是否开启调试模式，默认为False
        :param max_value: 像素最大值，默认为255
        """
        super().__init__(arch, debug, max_value)
    
    @ti.kernel
    def dither_atkinson(self, img: ti.types.ndarray()):
        """
        Atkinson抖动算法实现。
        
        :param img: 输入图像数组
        """
        h, w = img.shape
        for i, j in ti.ndrange(h, w):
            oldpixel = img[i, j]
            
            # 找到最接近的调色板颜色
            min_distance = self.MAX
            newpixel = 0
            for c in ti.static(self.PALETTE):
                distance = abs(oldpixel - c)
                if distance < min_distance:
                    min_distance = distance
                    newpixel = c
                    
            img[i, j] = newpixel
            quant_err = oldpixel - newpixel
            
            # Atkinson抖动权重分布 (1/8)
            err_part = quant_err >> 3
            
            # 周围像素的误差扩散
            if j + 1 < w: img[i, j + 1] += err_part
            if j + 2 < w: img[i, j + 2] += err_part
            
            if i + 1 < h:
                if j - 1 >= 0: img[i + 1, j - 1] += err_part
                img[i + 1, j] += err_part
                if j + 1 < w: img[i + 1, j + 1] += err_part
            
            if i + 2 < h:
                img[i + 2, j] += err_part
    
    # 添加Atkinson抖动的内核实现
    @ti.kernel
    def apply_color_atkinson_kernel(self, 
                                 img_in: ti.types.ndarray(),
                                 img_out: ti.types.ndarray(),
                                 palette: ti.types.ndarray(),
                                 intermediate: ti.types.ndarray()):
        """
        使用Taichi内核实现的彩色Atkinson抖动算法
        
        :param img_in: 输入彩色图像
        :param img_out: 输出彩色图像
        :param palette: 颜色调色板，形状为(n, 3)
        :param intermediate: 中间处理图像 (用于误差扩散)
        """
        h, w = img_in.shape[0], img_in.shape[1]
        n_colors = palette.shape[0]
        
        # 首先复制输入图像到中间缓冲区
        for i, j, c in ti.ndrange(h, w, 3):
            intermediate[i, j, c] = img_in[i, j, c]
        
        # 逐像素处理
        for i, j in ti.ndrange(h, w):
            # 获取当前像素值
            r = intermediate[i, j, 0]
            g = intermediate[i, j, 1]
            b = intermediate[i, j, 2]
            
            # 找到最接近的调色板颜色
            min_dist = 1e10
            best_idx = 0
            
            for k in range(n_colors):
                pr = palette[k, 0]
                pg = palette[k, 1]
                pb = palette[k, 2]
                
                # 计算颜色距离
                dist = (r-pr)*(r-pr) + (g-pg)*(g-pg) + (b-pb)*(b-pb)
                
                if dist < min_dist:
                    min_dist = dist
                    best_idx = k
            
            # 输出最佳匹配的颜色
            img_out[i, j, 0] = palette[best_idx, 0]
            img_out[i, j, 1] = palette[best_idx, 1]
            img_out[i, j, 2] = palette[best_idx, 2]
            
            # 计算误差
            err_r = r - palette[best_idx, 0]
            err_g = g - palette[best_idx, 1]
            err_b = b - palette[best_idx, 2]
            
            # Atkinson误差扩散 (1/8)
            err_r_part = err_r // 8
            err_g_part = err_g // 8
            err_b_part = err_b // 8
            
            # 扩散误差到周围像素
            if j + 1 < w:
                intermediate[i, j + 1, 0] = ti.min(ti.max(intermediate[i, j + 1, 0] + err_r_part, 0), self.MAX)
                intermediate[i, j + 1, 1] = ti.min(ti.max(intermediate[i, j + 1, 1] + err_g_part, 0), self.MAX)
                intermediate[i, j + 1, 2] = ti.min(ti.max(intermediate[i, j + 1, 2] + err_b_part, 0), self.MAX)
                
            if j + 2 < w:
                intermediate[i, j + 2, 0] = ti.min(ti.max(intermediate[i, j + 2, 0] + err_r_part, 0), self.MAX)
                intermediate[i, j + 2, 1] = ti.min(ti.max(intermediate[i, j + 2, 1] + err_g_part, 0), self.MAX)
                intermediate[i, j + 2, 2] = ti.min(ti.max(intermediate[i, j + 2, 2] + err_b_part, 0), self.MAX)
                
            if i + 1 < h:
                if j - 1 >= 0:
                    intermediate[i + 1, j - 1, 0] = ti.min(ti.max(intermediate[i + 1, j - 1, 0] + err_r_part, 0), self.MAX)
                    intermediate[i + 1, j - 1, 1] = ti.min(ti.max(intermediate[i + 1, j - 1, 1] + err_g_part, 0), self.MAX)
                    intermediate[i + 1, j - 1, 2] = ti.min(ti.max(intermediate[i + 1, j - 1, 2] + err_b_part, 0), self.MAX)
                    
                intermediate[i + 1, j, 0] = ti.min(ti.max(intermediate[i + 1, j, 0] + err_r_part, 0), self.MAX)
                intermediate[i + 1, j, 1] = ti.min(ti.max(intermediate[i + 1, j, 1] + err_g_part, 0), self.MAX)
                intermediate[i + 1, j, 2] = ti.min(ti.max(intermediate[i + 1, j, 2] + err_b_part, 0), self.MAX)
                
                if j + 1 < w:
                    intermediate[i + 1, j + 1, 0] = ti.min(ti.max(intermediate[i + 1, j + 1, 0] + err_r_part, 0), self.MAX)
                    intermediate[i + 1, j + 1, 1] = ti.min(ti.max(intermediate[i + 1, j + 1, 1] + err_g_part, 0), self.MAX)
                    intermediate[i + 1, j + 1, 2] = ti.min(ti.max(intermediate[i + 1, j + 1, 2] + err_b_part, 0), self.MAX)
                    
            if i + 2 < h:
                intermediate[i + 2, j, 0] = ti.min(ti.max(intermediate[i + 2, j, 0] + err_r_part, 0), self.MAX)
                intermediate[i + 2, j, 1] = ti.min(ti.max(intermediate[i + 2, j, 1] + err_g_part, 0), self.MAX)
                intermediate[i + 2, j, 2] = ti.min(ti.max(intermediate[i + 2, j, 2] + err_b_part, 0), self.MAX)
    
    # 修改后的彩色Atkinson抖动实现，使用上面的内核函数
    def apply_color_atkinson(self, img, color_palette):
        """
        应用彩色Atkinson抖动算法。
        
        :param img: 输入彩色图像numpy数组 (h, w, 3)
        :param color_palette: 颜色调色板 [(r,g,b), ...]
        :return: 处理后的彩色图像
        """
        h, w, c = img.shape
        
        # 转换为numpy数组处理
        img_np = img.copy()
        
        # 将调色板转换为numpy数组
        palette_np = np.array(color_palette, dtype=np.int32)
        
        # 创建输出图像
        output_img = np.zeros_like(img_np)
        
        # 使用GPU加速的Taichi内核处理
        if self.arch == ti.gpu:
            # 转换为Taichi数组
            img_ti = ti.ndarray(shape=img_np.shape, dtype=ti.i32)
            img_ti.from_numpy(img_np.astype(np.int32))
            
            output_ti = ti.ndarray(shape=output_img.shape, dtype=ti.i32)
            
            # 中间缓冲区，用于误差扩散
            intermediate_ti = ti.ndarray(shape=img_np.shape, dtype=ti.i32)
            
            palette_ti = ti.ndarray(shape=palette_np.shape, dtype=ti.i32)
            palette_ti.from_numpy(palette_np)
            
            # 执行GPU加速的内核函数
            self.apply_color_atkinson_kernel(img_ti, output_ti, palette_ti, intermediate_ti)
            
            # 确保最终图像只使用调色板中的颜色
            # 使用Taichi内核加速颜色映射
            self.map_to_color_palette(output_ti, palette_ti)
            
            # 转回numpy数组
            output_img = output_ti.to_numpy()
        else:
            # 回退到CPU处理方式
            # 使用更高效的颜色匹配算法
            palette_tree = cKDTree(palette_np)
            
            # 显示进度
            for i in tqdm(range(h), desc="处理图像", unit="行"):
                for j in range(w):
                    # 获取当前像素
                    pixel = img_np[i, j].astype(np.int32)
                    
                    # 使用KD树查找最近颜色
                    _, idx = palette_tree.query(pixel)
                    best_color = palette_np[idx]
                    
                    # 应用最近的颜色
                    output_img[i, j] = best_color
                    
                    # 计算量化误差
                    quant_error = pixel - best_color
                    
                    # Atkinson误差扩散 (1/8)
                    err_part = quant_error // 8
                    
                    # 周围像素的误差扩散
                    if j + 1 < w: 
                        img_np[i, j + 1] = np.clip(img_np[i, j + 1] + err_part, 0, self.MAX)
                    if j + 2 < w: 
                        img_np[i, j + 2] = np.clip(img_np[i, j + 2] + err_part, 0, self.MAX)
                    
                    if i + 1 < h:
                        if j - 1 >= 0: 
                            img_np[i + 1, j - 1] = np.clip(img_np[i + 1, j - 1] + err_part, 0, self.MAX)
                        img_np[i + 1, j] = np.clip(img_np[i + 1, j] + err_part, 0, self.MAX)
                        if j + 1 < w: 
                            img_np[i + 1, j + 1] = np.clip(img_np[i + 1, j + 1] + err_part, 0, self.MAX)
                    
                    if i + 2 < h:
                        img_np[i + 2, j] = np.clip(img_np[i + 2, j] + err_part, 0, self.MAX)
            
            # 确保最终图像只使用调色板中的颜色（使用向量化操作加速）
            output_img_reshaped = output_img.reshape(-1, 3)
            distances = np.sum((output_img_reshaped[:, np.newaxis, :] - palette_np[np.newaxis, :, :]) ** 2, axis=2)
            indices = np.argmin(distances, axis=1)
            output_img = palette_np[indices].reshape(output_img.shape)
        
        return output_img.astype(np.uint8)
    
    def process_image(self, img_path: str, num_colors: int = 2, use_color: bool = False, 
                      output_width: int = None, output_height: int = None):
        """
        处理图像并应用Atkinson抖动算法。
        
        :param img_path: 图像文件路径
        :param num_colors: 色值数量，默认为2（黑白）
        :param use_color: 是否处理彩色图像，默认为False
        :param output_width: 输出图像宽度，如果为None则保持原始宽度
        :param output_height: 输出图像高度，如果为None则保持原始高度
        :return: 处理后的图像
        """
        # 读取图像
        original_img = cv2.imread(img_path)
        
        if original_img is None:
            raise ValueError(f"无法读取图像: {img_path}")
            
        # 调整图像大小（如果指定了输出尺寸）
        if output_width is not None or output_height is not None:
            original_img = self.resize_image(original_img, output_width, output_height)
        
        if use_color:
            # 彩色图像处理
            self.color_mode = 'color'
            
            # 从图像中提取调色板
            color_palette = self.generate_color_palette_from_image(original_img, num_colors)
            
            if self.DEBUG:
                print(f"已从图像中提取 {num_colors} 种颜色: {color_palette}")
            
            # 制作一个副本用于处理
            img = original_img.copy()
            
            # 应用彩色Atkinson抖动
            img = self.apply_color_atkinson(img, color_palette)
            
            return img
            
        else:
            # 灰度图像处理
            self.color_mode = 'grayscale'
            
            # 生成灰度调色板
            self.generate_palette(num_colors)
            
            # 转换为灰度图
            img = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY)
            img = img.astype(np.int32)  # 转换为 int32 类型

            if self.arch == ti.gpu:
                # 转换为Taichi数组
                img_ti = ti.ndarray(shape=img.shape, dtype=ti.u8)
                img_ti.from_numpy(img)
            
                # 应用Atkinson抖动
                self.dither_atkinson(img_ti)
            else:
                self.dither_atkinson(img)

            # 转换回uint8
            img = np.clip(img, 0, self.MAX)
            img = img.astype(np.uint8)
            
            # 确保所有像素值严格使用调色板中的颜色
            palette_np = np.array(self.PALETTE, dtype=np.uint8)
            
            # 使用Taichi内核加速颜色映射
            if self.arch == ti.gpu:
                palette_ti = ti.ndarray(shape=palette_np.shape, dtype=ti.u8)
                palette_ti.from_numpy(palette_np)
                
                # 使用Taichi内核加速映射
                self.map_to_palette(img_ti, palette_ti)
                
                # 转回numpy数组
                img = img_ti.to_numpy()
            else:
                # 使用向量化操作加速CPU处理
                img_flat = img.flatten()
                idx = np.argmin(np.abs(img_flat.reshape(-1, 1) - palette_np.reshape(1, -1)), axis=1)
                img_flat[:] = palette_np[idx]
                img = img_flat.reshape(img.shape)
            
            return img 